#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
故障恢复管理器
功能：高级故障恢复管理，包括依赖关系处理、恢复策略、历史记录等
作者：Ansible自动生成
"""

import os
import sys
import json
import yaml
import time
import logging
import argparse
import subprocess
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional, Tuple
from pathlib import Path


class FailoverManager:
    """故障恢复管理器主类"""
    
    def __init__(self, config_file: str = '{{ failover_config_dir }}/config.yml'):
        self.config_file = config_file
        self.config = self.load_config()
        self.logger = self.setup_logging()
        self.history_file = self.config['global'].get('history_file', '/var/log/failover/history.json')
        self.dependency_graph = self.build_dependency_graph()
        
    def load_config(self) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"Error loading config file {self.config_file}: {e}")
            sys.exit(1)
    
    def setup_logging(self) -> logging.Logger:
        """设置日志记录"""
        log_file = self.config['global']['log_file']
        log_level = getattr(logging, self.config['global']['log_level'].upper())
        
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        
        logging.basicConfig(
            level=log_level,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler(sys.stdout)
            ]
        )
        
        return logging.getLogger('failover_manager')
    
    def build_dependency_graph(self) -> Dict[str, List[str]]:
        """构建服务依赖关系图"""
        graph = {}
        services = self.config.get('services', [])
        
        for service in services:
            service_name = service['name']
            dependencies = service.get('dependencies', [])
            graph[service_name] = dependencies
        
        return graph
    
    def get_service_status(self, service_name: str) -> Dict[str, Any]:
        """获取服务状态"""
        try:
            # 检查systemd服务状态
            result = subprocess.run(
                ['systemctl', 'show', service_name, '--property=ActiveState,SubState,LoadState'],
                capture_output=True,
                text=True
            )
            
            if result.returncode == 0:
                status = {}
                for line in result.stdout.strip().split('\n'):
                    if '=' in line:
                        key, value = line.split('=', 1)
                        status[key] = value
                
                return {
                    'active': status.get('ActiveState', 'unknown'),
                    'sub_state': status.get('SubState', 'unknown'),
                    'load_state': status.get('LoadState', 'unknown'),
                    'timestamp': datetime.now().isoformat()
                }
            else:
                return {
                    'active': 'unknown',
                    'sub_state': 'unknown',
                    'load_state': 'unknown',
                    'error': result.stderr,
                    'timestamp': datetime.now().isoformat()
                }
                
        except Exception as e:
            self.logger.error(f"Error getting status for {service_name}: {e}")
            return {
                'active': 'error',
                'error': str(e),
                'timestamp': datetime.now().isoformat()
            }
    
    def restart_service_with_dependencies(self, service_name: str, force: bool = False) -> bool:
        """重启服务及其依赖"""
        self.logger.info(f"Starting restart process for {service_name}")
        
        # 获取重启顺序
        restart_order = self.get_restart_order(service_name)
        self.logger.info(f"Restart order: {' -> '.join(restart_order)}")
        
        success = True
        failed_services = []
        
        for svc in restart_order:
            try:
                self.logger.info(f"Restarting service: {svc}")
                
                # 停止服务
                stop_result = subprocess.run(
                    ['systemctl', 'stop', svc],
                    capture_output=True,
                    text=True,
                    timeout=30
                )
                
                if stop_result.returncode != 0:
                    self.logger.warning(f"Failed to stop {svc}: {stop_result.stderr}")
                    if not force:
                        failed_services.append(svc)
                        success = False
                        continue
                
                # 等待一段时间
                time.sleep(2)
                
                # 启动服务
                start_result = subprocess.run(
                    ['systemctl', 'start', svc],
                    capture_output=True,
                    text=True,
                    timeout=60
                )
                
                if start_result.returncode != 0:
                    self.logger.error(f"Failed to start {svc}: {start_result.stderr}")
                    failed_services.append(svc)
                    success = False
                    if not force:
                        break
                else:
                    self.logger.info(f"Successfully restarted {svc}")
                    
                    # 等待服务稳定
                    time.sleep(5)
                    
                    # 验证服务状态
                    status = self.get_service_status(svc)
                    if status['active'] != 'active':
                        self.logger.warning(f"Service {svc} is not active after restart: {status['active']}")
                        if not force:
                            failed_services.append(svc)
                            success = False
                            break
                
            except subprocess.TimeoutExpired:
                self.logger.error(f"Timeout while restarting {svc}")
                failed_services.append(svc)
                success = False
                if not force:
                    break
            except Exception as e:
                self.logger.error(f"Error restarting {svc}: {e}")
                failed_services.append(svc)
                success = False
                if not force:
                    break
        
        # 记录重启历史
        self.record_restart_history(service_name, restart_order, success, failed_services)
        
        if success:
            self.logger.info(f"Successfully restarted {service_name} and dependencies")
        else:
            self.logger.error(f"Failed to restart some services: {failed_services}")
        
        return success
    
    def get_restart_order(self, service_name: str) -> List[str]:
        """获取服务重启顺序（拓扑排序）"""
        visited = set()
        temp_visited = set()
        order = []
        
        def dfs(svc: str):
            if svc in temp_visited:
                raise ValueError(f"Circular dependency detected involving {svc}")
            if svc in visited:
                return
            
            temp_visited.add(svc)
            
            # 访问依赖的服务
            for dep in self.dependency_graph.get(svc, []):
                dfs(dep)
            
            temp_visited.remove(svc)
            visited.add(svc)
            order.append(svc)
        
        try:
            dfs(service_name)
            return order
        except ValueError as e:
            self.logger.error(f"Dependency error: {e}")
            return [service_name]  # 如果有循环依赖，只重启目标服务
    
    def health_recovery(self, service_name: str) -> bool:
        """健康恢复：尝试多种恢复策略"""
        self.logger.info(f"Starting health recovery for {service_name}")
        
        service_config = self.get_service_config(service_name)
        if not service_config:
            self.logger.error(f"Service {service_name} not found in configuration")
            return False
        
        recovery_config = service_config.get('recovery', {})
        strategy = recovery_config.get('strategy', 'progressive')
        
        if strategy == 'progressive':
            return self.progressive_recovery(service_name, service_config)
        elif strategy == 'immediate':
            return self.immediate_recovery(service_name, service_config)
        elif strategy == 'manual':
            self.logger.info(f"Manual recovery required for {service_name}")
            return False
        else:
            self.logger.warning(f"Unknown recovery strategy: {strategy}")
            return self.progressive_recovery(service_name, service_config)
    
    def progressive_recovery(self, service_name: str, service_config: Dict[str, Any]) -> bool:
        """渐进式恢复：从轻量级到重量级的恢复步骤"""
        self.logger.info(f"Starting progressive recovery for {service_name}")
        
        # 步骤1：简单重启
        self.logger.info("Step 1: Simple restart")
        if self.simple_restart(service_name):
            if self.verify_service_health(service_name, service_config):
                self.logger.info(f"Service {service_name} recovered with simple restart")
                return True
        
        # 步骤2：强制重启
        self.logger.info("Step 2: Force restart")
        if self.force_restart(service_name):
            if self.verify_service_health(service_name, service_config):
                self.logger.info(f"Service {service_name} recovered with force restart")
                return True
        
        # 步骤3：重启依赖服务
        self.logger.info("Step 3: Restart with dependencies")
        if self.restart_service_with_dependencies(service_name):
            if self.verify_service_health(service_name, service_config):
                self.logger.info(f"Service {service_name} recovered with dependency restart")
                return True
        
        # 步骤4：执行自定义恢复脚本
        recovery_script = service_config.get('recovery', {}).get('custom_script')
        if recovery_script:
            self.logger.info("Step 4: Custom recovery script")
            if self.run_custom_recovery(service_name, recovery_script):
                if self.verify_service_health(service_name, service_config):
                    self.logger.info(f"Service {service_name} recovered with custom script")
                    return True
        
        self.logger.error(f"All recovery steps failed for {service_name}")
        return False
    
    def immediate_recovery(self, service_name: str, service_config: Dict[str, Any]) -> bool:
        """立即恢复：直接重启服务和依赖"""
        self.logger.info(f"Starting immediate recovery for {service_name}")
        return self.restart_service_with_dependencies(service_name, force=True)
    
    def simple_restart(self, service_name: str) -> bool:
        """简单重启服务"""
        try:
            result = subprocess.run(
                ['systemctl', 'restart', service_name],
                capture_output=True,
                text=True,
                timeout=60
            )
            
            if result.returncode == 0:
                self.logger.info(f"Successfully restarted {service_name}")
                return True
            else:
                self.logger.warning(f"Failed to restart {service_name}: {result.stderr}")
                return False
                
        except Exception as e:
            self.logger.error(f"Error during simple restart of {service_name}: {e}")
            return False
    
    def force_restart(self, service_name: str) -> bool:
        """强制重启服务"""
        try:
            # 强制停止
            subprocess.run(['systemctl', 'kill', service_name], capture_output=True)
            time.sleep(2)
            
            # 重启
            result = subprocess.run(
                ['systemctl', 'restart', service_name],
                capture_output=True,
                text=True,
                timeout=60
            )
            
            if result.returncode == 0:
                self.logger.info(f"Successfully force restarted {service_name}")
                return True
            else:
                self.logger.warning(f"Failed to force restart {service_name}: {result.stderr}")
                return False
                
        except Exception as e:
            self.logger.error(f"Error during force restart of {service_name}: {e}")
            return False
    
    def run_custom_recovery(self, service_name: str, script: str) -> bool:
        """运行自定义恢复脚本"""
        try:
            self.logger.info(f"Running custom recovery script for {service_name}")
            
            result = subprocess.run(
                script,
                shell=True,
                capture_output=True,
                text=True,
                timeout=300,
                env={**os.environ, 'SERVICE_NAME': service_name}
            )
            
            if result.returncode == 0:
                self.logger.info(f"Custom recovery script succeeded for {service_name}")
                return True
            else:
                self.logger.error(f"Custom recovery script failed for {service_name}: {result.stderr}")
                return False
                
        except subprocess.TimeoutExpired:
            self.logger.error(f"Custom recovery script timed out for {service_name}")
            return False
        except Exception as e:
            self.logger.error(f"Error running custom recovery script for {service_name}: {e}")
            return False
    
    def verify_service_health(self, service_name: str, service_config: Dict[str, Any]) -> bool:
        """验证服务健康状态"""
        # 等待服务稳定
        time.sleep(10)
        
        # 检查systemd状态
        status = self.get_service_status(service_name)
        if status['active'] != 'active':
            self.logger.warning(f"Service {service_name} is not active: {status['active']}")
            return False
        
        # 运行健康检查（如果配置了）
        health_check = service_config.get('health_check')
        if health_check:
            try:
                if health_check.get('type') == 'http':
                    import requests
                    url = health_check['url']
                    response = requests.get(url, timeout=10)
                    if response.status_code != health_check.get('expected_status', 200):
                        self.logger.warning(f"Health check failed for {service_name}: HTTP {response.status_code}")
                        return False
                elif health_check.get('type') == 'tcp':
                    import socket
                    host = health_check['host']
                    port = health_check['port']
                    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                    sock.settimeout(5)
                    result = sock.connect_ex((host, port))
                    sock.close()
                    if result != 0:
                        self.logger.warning(f"Health check failed for {service_name}: TCP connection failed")
                        return False
                elif health_check.get('type') == 'script':
                    script = health_check['script']
                    result = subprocess.run(
                        script,
                        shell=True,
                        capture_output=True,
                        text=True,
                        timeout=30
                    )
                    if result.returncode != 0:
                        self.logger.warning(f"Health check failed for {service_name}: Script returned {result.returncode}")
                        return False
            except Exception as e:
                self.logger.error(f"Error during health check for {service_name}: {e}")
                return False
        
        self.logger.info(f"Service {service_name} health verification passed")
        return True
    
    def get_service_config(self, service_name: str) -> Optional[Dict[str, Any]]:
        """获取服务配置"""
        services = self.config.get('services', [])
        for service in services:
            if service['name'] == service_name:
                return service
        return None
    
    def record_restart_history(self, service_name: str, restart_order: List[str], 
                             success: bool, failed_services: List[str]):
        """记录重启历史"""
        try:
            history_entry = {
                'timestamp': datetime.now().isoformat(),
                'service': service_name,
                'restart_order': restart_order,
                'success': success,
                'failed_services': failed_services,
                'duration': 0  # 可以添加计时功能
            }
            
            # 读取现有历史
            history = []
            if os.path.exists(self.history_file):
                try:
                    with open(self.history_file, 'r') as f:
                        history = json.load(f)
                except Exception:
                    pass
            
            # 添加新记录
            history.append(history_entry)
            
            # 保持最近1000条记录
            if len(history) > 1000:
                history = history[-1000:]
            
            # 保存历史
            os.makedirs(os.path.dirname(self.history_file), exist_ok=True)
            with open(self.history_file, 'w') as f:
                json.dump(history, f, indent=2)
                
        except Exception as e:
            self.logger.error(f"Error recording restart history: {e}")
    
    def get_restart_history(self, service_name: str = None, days: int = 7) -> List[Dict[str, Any]]:
        """获取重启历史"""
        try:
            if not os.path.exists(self.history_file):
                return []
            
            with open(self.history_file, 'r') as f:
                history = json.load(f)
            
            # 过滤时间范围
            cutoff_time = datetime.now() - timedelta(days=days)
            filtered_history = [
                entry for entry in history
                if datetime.fromisoformat(entry['timestamp']) > cutoff_time
            ]
            
            # 过滤服务名称
            if service_name:
                filtered_history = [
                    entry for entry in filtered_history
                    if entry['service'] == service_name
                ]
            
            return filtered_history
            
        except Exception as e:
            self.logger.error(f"Error getting restart history: {e}")
            return []
    
    def generate_report(self, days: int = 7) -> str:
        """生成故障恢复报告"""
        history = self.get_restart_history(days=days)
        
        if not history:
            return "No restart history found in the specified time period."
        
        # 统计信息
        total_restarts = len(history)
        successful_restarts = sum(1 for entry in history if entry['success'])
        failed_restarts = total_restarts - successful_restarts
        
        # 按服务统计
        service_stats = {}
        for entry in history:
            service = entry['service']
            if service not in service_stats:
                service_stats[service] = {'total': 0, 'success': 0, 'failed': 0}
            
            service_stats[service]['total'] += 1
            if entry['success']:
                service_stats[service]['success'] += 1
            else:
                service_stats[service]['failed'] += 1
        
        # 生成报告
        report = f"""
故障恢复报告 (最近 {days} 天)
生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

=== 总体统计 ===
总重启次数: {total_restarts}
成功重启: {successful_restarts}
失败重启: {failed_restarts}
成功率: {(successful_restarts/total_restarts*100):.1f}%

=== 服务统计 ===
"""
        
        for service, stats in service_stats.items():
            success_rate = (stats['success'] / stats['total'] * 100) if stats['total'] > 0 else 0
            report += f"""
服务: {service}
  总重启: {stats['total']}
  成功: {stats['success']}
  失败: {stats['failed']}
  成功率: {success_rate:.1f}%
"""
        
        # 最近的失败记录
        recent_failures = [entry for entry in history[-10:] if not entry['success']]
        if recent_failures:
            report += "\n=== 最近失败记录 ===\n"
            for entry in recent_failures:
                report += f"""
时间: {entry['timestamp']}
服务: {entry['service']}
失败服务: {', '.join(entry['failed_services'])}
"""
        
        return report


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='故障恢复管理器')
    parser.add_argument('action', choices=['restart', 'recover', 'status', 'history', 'report'],
                       help='执行的操作')
    parser.add_argument('--service', '-s', help='服务名称')
    parser.add_argument('--force', '-f', action='store_true', help='强制执行')
    parser.add_argument('--days', '-d', type=int, default=7, help='历史记录天数')
    parser.add_argument('--config', '-c', default='{{ failover_config_dir }}/config.yml', help='配置文件路径')
    
    args = parser.parse_args()
    
    try:
        manager = FailoverManager(args.config)
        
        if args.action == 'restart':
            if not args.service:
                print("Error: --service is required for restart action")
                sys.exit(1)
            
            success = manager.restart_service_with_dependencies(args.service, args.force)
            sys.exit(0 if success else 1)
            
        elif args.action == 'recover':
            if not args.service:
                print("Error: --service is required for recover action")
                sys.exit(1)
            
            success = manager.health_recovery(args.service)
            sys.exit(0 if success else 1)
            
        elif args.action == 'status':
            if args.service:
                status = manager.get_service_status(args.service)
                print(json.dumps(status, indent=2))
            else:
                services = manager.config.get('services', [])
                for service in services:
                    status = manager.get_service_status(service['name'])
                    print(f"{service['name']}: {status['active']}")
                    
        elif args.action == 'history':
            history = manager.get_restart_history(args.service, args.days)
            print(json.dumps(history, indent=2))
            
        elif args.action == 'report':
            report = manager.generate_report(args.days)
            print(report)
            
    except KeyboardInterrupt:
        print("\nOperation interrupted by user")
        sys.exit(0)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()